import asyncio
import json
import subprocess
import sys

class SimpleMCPClient:
    def __init__(self, server_process):
        self.server_process = server_process
        self.request_id = 1

    def write_message(self, message: dict):
        """写入MCP格式的消息"""
        content = json.dumps(message)
        message_str = f"Content-Length: {len(content)}\r\n\r\n{content}"
        self.server_process.stdin.write(message_str.encode('utf-8'))
        self.server_process.stdin.flush()

    def read_message(self):
        """读取MCP格式的消息"""
        # 读取Content-Length头
        line = self.server_process.stdout.readline().decode('utf-8').strip()
        if not line:
            return None

        if line.startswith('Content-Length:'):
            # 提取内容长度
            content_length = int(line.split(':')[1].strip())

            # 读取空行
            self.server_process.stdout.readline()

            # 读取JSON内容
            content = self.server_process.stdout.read(content_length).decode('utf-8')
            return json.loads(content)

        return None

    async def send_request(self, method: str, params: dict = None) -> dict:
        if params is None:
            params = {}

        request = {
            "jsonrpc": "2.0",
            "id": self.request_id,
            "method": method,
            "params": params
        }
        self.request_id += 1

        self.write_message(request)
        return self.read_message()

async def main():
    # 启动服务端进程
    server_process = subprocess.Popen(
        [sys.executable, "SimpleMCPServer.py"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=sys.stderr  # 将stderr重定向到控制台以便查看错误信息
    )

    client = SimpleMCPClient(server_process)

    try:
        # 1. 初始化连接
        print("Initializing connection...")
        init_response = await client.send_request("initialize", {})
        print(f"Initialization response: {json.dumps(init_response, indent=2)}")

        # 2. 列出可用资源
        print("\nListing resources...")
        list_response = await client.send_request("resources/list")
        resources = list_response.get('result', {}).get('resources', [])
        print(f"Available resources: {[r['name'] for r in resources]}")

        # 3. 读取资源内容
        print("\nReading resources...")
        for resource in resources:
            read_response = await client.send_request("resources/read", {"uri": resource["uri"]})
            content = read_response.get('result', {}).get('contents', [{}])[0].get('content', '')
            print(f"{resource['name']}: {content}")

        # 4. 关闭连接
        print("\nShutting down...")
        shutdown_response = await client.send_request("shutdown")
        print("Shutdown complete")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        server_process.terminate()

if __name__ == "__main__":
    asyncio.run(main())